"""
Created on Mar 1, 2020
Pytorch Implementation of LightGCN in
Xiangnan He et al. LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation

@author: Jianbai Ye (gusye@mail.ustc.edu.cn)

Define models here
"""
import torch
from torch import nn


def cprint(words : str):
    print(f"\033[0;30;43m{words}\033[0m")


class BasicModel(nn.Module):    
    def __init__(self, config, dataset):
        super(BasicModel, self).__init__()
    
    def getUsersRating(self, users):
        raise NotImplementedError
    
class PairWiseModel(BasicModel):
    def __init__(self):
        super(PairWiseModel, self).__init__()
    def bpr_loss(self, users, pos, neg, numAllPos):
        """
        Parameters:
            users: users list 
            pos: positive items for corresponding users
            neg: negative items for corresponding users
        Return:
            (log-loss, l2-loss)
        """
        raise NotImplementedError
    
class PureMF(BasicModel):
    def __init__(self, 
                 config:dict, 
                 dataset):
        super(PureMF, self).__init__()
        self.num_users  = dataset.n_users
        self.num_items  = dataset.m_items
        self.latent_dim = config['latent_dim_rec']
        self.f = nn.Sigmoid()
        self.__init_weight()
        
    def __init_weight(self):
        self.embedding_user = torch.nn.Embedding(
            num_embeddings=self.num_users, embedding_dim=self.latent_dim)
        self.embedding_item = torch.nn.Embedding(
            num_embeddings=self.num_items, embedding_dim=self.latent_dim)
        print("using Normal distribution N(0,1) initialization for PureMF")
        
    def getUsersRating(self, users):
        users = users.long()
        users_emb = self.embedding_user(users)
        items_emb = self.embedding_item.weight
        scores = torch.matmul(users_emb, items_emb.t())
        return self.f(scores)
    
    def bpr_loss(self, users, pos, neg):
        users_emb = self.embedding_user(users.long())
        pos_emb   = self.embedding_item(pos.long())
        neg_emb   = self.embedding_item(neg.long())
        pos_scores= torch.sum(users_emb*pos_emb, dim=1)
        neg_scores= torch.sum(users_emb*neg_emb, dim=1)
        loss = torch.mean(nn.functional.softplus(neg_scores - pos_scores))
        reg_loss = (1/2)*(users_emb.norm(2).pow(2) + 
                          pos_emb.norm(2).pow(2) + 
                          neg_emb.norm(2).pow(2))/float(len(users))
        return loss, reg_loss
        
    def forward(self, users, items):
        users = users.long()
        items = items.long()
        users_emb = self.embedding_user(users)
        items_emb = self.embedding_item(items)
        scores = torch.sum(users_emb*items_emb, dim=1)
        return self.f(scores)

class LightGCN(BasicModel):
    def __init__(self, config, dataset):
        super(LightGCN, self).__init__(config, dataset)
        self.config = config
        self.dataset = dataset
        self.init_weight()

    def init_weight(self):
        self.num_users  = self.dataset.n_users
        self.num_items  = self.dataset.m_items
        self.latent_dim = self.config['latent_dim_rec']
        self.n_layers = self.config['lightGCN_n_layers']
        self.keep_prob = self.config['keep_prob']
        self.A_split = self.config['A_split']
        self.embedding_user = torch.nn.Embedding(
            num_embeddings=self.num_users, embedding_dim=self.latent_dim)
        self.embedding_item = torch.nn.Embedding(
            num_embeddings=self.num_items, embedding_dim=self.latent_dim)
        if self.config['pretrain'] == 0:
#             nn.init.xavier_uniform_(self.embedding_user.weight, gain=1)
#             nn.init.xavier_uniform_(self.embedding_item.weight, gain=1)
#             print('use xavier initilizer')
# random normal init seems to be a better choice when lightGCN actually don't use any non-linear activation function
            nn.init.normal_(self.embedding_user.weight, std=0.1)
            nn.init.normal_(self.embedding_item.weight, std=0.1)
            cprint('use NORMAL distribution initilizer')
        else:
            self.embedding_user.weight.data.copy_(torch.from_numpy(self.config['user_emb']))
            self.embedding_item.weight.data.copy_(torch.from_numpy(self.config['item_emb']))
            print('use pretarined data')
        self.f = nn.Sigmoid()
        self.Graph = self.dataset.getSparseGraph()
        print(f"lgn is already to go(dropout:{self.config['dropout']})")

        # print("save_txt")
    def __dropout_x(self, x, keep_prob):
        size = x.size()
        index = x.indices().t()
        values = x.values()
        random_index = torch.rand(len(values)) + keep_prob
        random_index = random_index.int().bool()
        index = index[random_index]
        values = values[random_index]/keep_prob
        g = torch.sparse.FloatTensor(index.t(), values, size)
        return g
    
    def __dropout(self, keep_prob):
        if self.A_split:
            graph = []
            for g in self.Graph:
                graph.append(self.__dropout_x(g, keep_prob))
        else:
            graph = self.__dropout_x(self.Graph, keep_prob)
        return graph
    
    def computer(self):
        """
        propagate methods for lightGCN
        """       
        users_emb = self.embedding_user.weight
        items_emb = self.embedding_item.weight
        all_emb = torch.cat([users_emb, items_emb])
        #   torch.split(all_emb , [self.num_users, self.num_items])
        embs = [all_emb]
        if self.config['dropout']:
            if self.training:
                print("droping")
                g_droped = self.__dropout(self.keep_prob)
            else:
                g_droped = self.Graph        
        else:
            g_droped = self.Graph    
        
        for layer in range(self.n_layers):
            if self.A_split:
                temp_emb = []
                for f in range(len(g_droped)):
                    temp_emb.append(torch.sparse.mm(g_droped[f], all_emb))
                side_emb = torch.cat(temp_emb, dim=0)
                all_emb = side_emb
            else:
                all_emb = torch.sparse.mm(g_droped, all_emb)
            embs.append(all_emb)
        embs = torch.stack(embs, dim=1)
        #print(embs.size())
        light_out = torch.mean(embs, dim=1)
        users, items = torch.split(light_out, [self.num_users, self.num_items])
        return users, items
    
    def getUsersRating(self, users):  # 按用户test
        all_users, all_items = self.computer()
        users_emb = all_users[users.long()]
        items_emb = all_items
        # users_emb = torch.sign(users_emb)  # 反向传播跳过sign函数
        # items_emb = torch.sign(items_emb)
        rating = self.f(torch.matmul(users_emb, items_emb.t()))
        return rating
    
    def getEmbedding(self, users, pos_items, neg_items):
        all_users, all_items = self.computer()
        users_emb = all_users[users]
        pos_emb = all_items[pos_items]
        neg_emb = all_items[neg_items]
        users_emb_ego = self.embedding_user(users)
        pos_emb_ego = self.embedding_item(pos_items)
        neg_emb_ego = self.embedding_item(neg_items)
        return users_emb, pos_emb, neg_emb, users_emb_ego, pos_emb_ego, neg_emb_ego

    def get_user_item_embedding(self, users, items):
        all_users, all_items = self.computer()
        users_emb = all_users[users]
        items_emb = all_items[items]
        users_emb_ego = self.embedding_user(users)
        items_emb_ego = self.embedding_item(items)
        return users_emb, items_emb, users_emb_ego, items_emb_ego
    
    # def bpr_loss(self, users, pos, neg):
    #     (users_emb, pos_emb, neg_emb,
    #     userEmb0,  posEmb0, negEmb0) = self.getEmbedding(users.long(), pos.long(), neg.long())
    #     reg_loss = (1/2)*(userEmb0.norm(2).pow(2) +
    #                      posEmb0.norm(2).pow(2)  +
    #                      negEmb0.norm(2).pow(2))/float(len(users))
    #     pos_scores = torch.mul(users_emb, pos_emb)
    #     pos_scores = torch.sum(pos_scores, dim=1)
    #     neg_scores = torch.mul(users_emb, neg_emb)
    #     neg_scores = torch.sum(neg_scores, dim=1)
    #
    #     loss = torch.mean(torchZz.nn.functional.softplus(neg_scores - pos_scores))
    #
    #     return loss, reg_loss

    def bpr_loss(self, users, pos, neg, numAllPos):
        tau = self.config["tau"]
        num_user_iteraction = numAllPos[users]
        (users_emb, pos_emb, neg_emb,
         userEmb0, posEmb0, negEmb0) = self.getEmbedding(users.long(), pos.long(), neg.long())
        reg_loss = (1 / 2) * (userEmb0.norm(2).pow(2) +
                              posEmb0.norm(2).pow(2) +
                              negEmb0.norm(2).pow(2)) / float(len(users))
        pos_scores = torch.mul(users_emb, pos_emb)
        pos_scores = torch.sum(pos_scores, dim=1)
        neg_scores = torch.mul(users_emb, neg_emb)
        neg_scores = torch.sum(neg_scores, dim=1)
        loss = 1-torch.mean(torch.div(torch.sigmoid((pos_scores - neg_scores)/tau), num_user_iteraction**2))
        return loss, reg_loss

    def forward(self, users, items):
        # compute embedding
        all_users, all_items = self.computer()
        # print('forward')
        #all_users, all_items = self.computer()
        users_emb = all_users[users]
        items_emb = all_items[items]
        inner_pro = torch.mul(users_emb, items_emb)
        gamma = torch.sum(inner_pro, dim=1)
        return gamma


class LightGCN_sign(LightGCN):
    def __init__(self, config, dataset):
        super(LightGCN_sign, self).__init__(config, dataset)
        self.config = config
        self.dataset = dataset
        self.init_weight()

    def getUsersRating(self, users):  # sign version
        all_users, all_items = self.computer()
        users_emb = all_users[users.long()]
        items_emb = all_items
        users_emb = torch.sign(users_emb)  # 反向传播跳过sign函数
        items_emb = torch.sign(items_emb)
        rating = self.f(torch.matmul(users_emb, items_emb.t()))
        return rating

    def forward(self, users, items):
        # compute embedding
        all_users, all_items = self.computer()
        # print('forward')
        # all_users, all_items = self.computer()
        users_emb = all_users[users]
        items_emb = all_items[items]

        users_emb = Sign.apply(users_emb)  # 反向传播跳过sign函数
        items_emb = Sign.apply(items_emb)
        inner_pro = torch.mul(users_emb, items_emb)
        gamma = torch.sum(inner_pro, dim=1)
        return gamma


class LightGCN_heaviside_2exp(LightGCN):
    def __init__(self, config, dataset):
        super(LightGCN_heaviside_2exp, self).__init__(config, dataset)
        self.two_exp = torch.pow(2, torch.arange(self.config["latent_dim_rec"])).to(self.config["device"])
        self.tmp = torch.tensor([0.5]).to(self.config["device"])
        self.fc_user = nn.Linear(64, 64)
        self.fc_item = nn.Linear(64, 64)

    def getUsersRating(self, users):  # sign version
        all_users, all_items = self.computer()
        users_emb = all_users[users.long()]
        items_emb = all_items
        users_emb = self.fc_user(users_emb)
        items_emb = self.fc_item(items_emb)
        users_emb = torch.heaviside(users_emb, self.tmp)  # 反向传播跳过sign函数
        items_emb = torch.heaviside(items_emb, self.tmp)
        users_emb = torch.mul(users_emb, self.two_exp)
        rating = self.f(torch.matmul(users_emb, items_emb.t()))
        return rating

    def forward(self, users, items):
        # compute embedding
        all_users, all_items = self.computer()
        users_emb = all_users[users]
        items_emb = all_items[items]
        users_emb = self.fc_user(users_emb)
        items_emb = self.fc_item(items_emb)
        # users_emb = users_emb * self.config["tanh_beta"]
        # items_emb = items_emb * self.config["tanh_beta"]
        users_emb = self.f(users_emb)  # 反向传播跳过sign函数
        items_emb = self.f(items_emb)
        users_emb = torch.mul(users_emb, self.two_exp)
        inner_pro = torch.mul(users_emb, items_emb)
        gamma = torch.sum(inner_pro, dim=1)
        return gamma


class LightGCN_fc(LightGCN):

    def __init__(self,
                 config: dict,
                 dataset):
        super(LightGCN_fc, self).__init__(config, dataset)
        self.config = config
        self.dataset = dataset
        # self.embedding_user_final = None
        # self.embedding_item_final = None
        self.init_weight()
        self.fc_user = nn.Linear(64, 64)
        self.fc_item = nn.Linear(64, 64)

    def getUsersRating(self, users):  # sign version
        all_users, all_items = self.computer()
        users_emb = all_users[users.long()]
        items_emb = all_items
        users_emb = self.fc_user(users_emb)
        items_emb = self.fc_item(items_emb)
        rating = self.f(torch.matmul(users_emb, items_emb.t()))
        return rating

    def forward(self, users, items):
        # compute embedding
        all_users, all_items = self.computer()
        # print('forward')
        # all_users, all_items = self.computer()
        users_emb = all_users[users]
        items_emb = all_items[items]
        users_emb = users_emb * self.config["tanh_beta"]
        items_emb = items_emb * self.config["tanh_beta"]
        users_emb = Sign_tanh.apply(users_emb)  # 反向传播跳过sign函数
        items_emb = Sign_tanh.apply(items_emb)
        inner_pro = torch.mul(users_emb, items_emb)
        gamma = torch.sum(inner_pro, dim=1)
        return gamma


class LightGCN_fc_sign_tanh(LightGCN):

    def __init__(self,
                 config: dict,
                 dataset):
        super(LightGCN_fc_sign_tanh, self).__init__(config, dataset)
        self.config = config
        self.dataset = dataset
        # self.embedding_user_final = None
        # self.embedding_item_final = None
        self.init_weight()
        self.fc_user = nn.Linear(64, 64)
        self.fc_item = nn.Linear(64, 64)

    def getUsersRating(self, users):  # sign version
        all_users, all_items = self.computer()
        users_emb = all_users[users.long()]
        items_emb = all_items
        users_emb = self.fc_user(users_emb)
        items_emb = self.fc_item(items_emb)
        users_emb = torch.sign(users_emb)  # 反向传播跳过sign函数
        items_emb = torch.sign(items_emb)
        rating = self.f(torch.matmul(users_emb, items_emb.t()))
        return rating

    def forward(self, users, items):
        # compute embedding
        all_users, all_items = self.computer()
        # print('forward')
        # all_users, all_items = self.computer()
        users_emb = all_users[users]
        items_emb = all_items[items]
        users_emb = self.fc_user(users_emb)
        items_emb = self.fc_item(items_emb)
        users_emb = users_emb * self.config["tanh_beta"]
        items_emb = items_emb * self.config["tanh_beta"]
        users_emb = Sign_tanh.apply(users_emb)  # 反向传播跳过sign函数
        items_emb = Sign_tanh.apply(items_emb)
        inner_pro = torch.mul(users_emb, items_emb)
        gamma = torch.sum(inner_pro, dim=1)
        return gamma


class LightGCN_sign_tanh(LightGCN):

    def __init__(self,
                 config: dict,
                 dataset):
        super(LightGCN_sign_tanh, self).__init__(config, dataset)
        self.config = config
        self.dataset = dataset
        self.init_weight()

    def getUsersRating(self, users):  # sign version
        all_users, all_items = self.computer()
        users_emb = all_users[users.long()]
        items_emb = all_items
        users_emb = torch.sign(users_emb)  # 反向传播跳过sign函数
        items_emb = torch.sign(items_emb)
        rating = self.f(torch.matmul(users_emb, items_emb.t()))
        return rating

    def forward(self, users, items):
        # compute embedding
        all_users, all_items = self.computer()
        users_emb = all_users[users]
        items_emb = all_items[items]
        users_emb = users_emb * self.config["tanh_beta"]
        items_emb = items_emb * self.config["tanh_beta"]
        users_emb = Sign_tanh.apply(users_emb)  # 反向传播跳过sign函数
        items_emb = Sign_tanh.apply(items_emb)
        inner_pro = torch.mul(users_emb, items_emb)
        gamma = torch.sum(inner_pro, dim=1)
        return gamma

class Sign(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return torch.sign(x)
    @staticmethod
    def backward(ctx, g):
        return g


class Sign_tanh(torch.autograd.Function):
    '''
    使用：
    output = sign_tanh.apply(input)
    '''
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.sign(input)
    @staticmethod
    def backward(ctx, grad_output):
        input = ctx.saved_variables
        # tanh(x)的导数 = 1 - (tanh(x))^2
        return (1 - torch.tanh(input)**2)*grad_output

# mf, lgn, lgn_sign, lgn_heaviside_2exp, lgn_sign_tanh, lgn_fc_sign_tanh
all_models = {
            'mf': PureMF,
            'lgn': LightGCN,
            "lgn_sign": LightGCN_sign,
            "lgn_heaviside_2exp": LightGCN_heaviside_2exp,
            "lgn_sign_tanh": LightGCN_sign_tanh,
            "lgn_fc_sign_tanh": LightGCN_fc_sign_tanh
        }


def check_model(model_name):
    if model_name not in all_models:
        raise NotImplementedError(f"Haven't supported {model_name} yet!, try {all_models.keys()}")


def getModel(model_name, config, dataset):
    if model_name in all_models:
        return all_models[model_name](config, dataset)
    else:
        raise NotImplementedError
